/**
 * https://leetcode.cn/submissions/detail/572467591/
 * 2385. 感染二叉树需要的总时间
 * Medium, 黄伟杰 2024.10.14
 * 树形DP
 */

#include <bits/stdc++.h>
using namespace std;

class Solution
{
    int ans = 0, start;
    pair<int, bool> dfs(TreeNode *node)
    {
        if (node == nullptr)
            return {0, false};
        auto [l_len, l_tag] = dfs(node->left);
        auto [r_len, r_tag] = dfs(node->right);
        if (node->val == start) // 计算start节点作为根节点的深度
        {
            ans = max(l_len, r_len);
            return {1, true}; // 当为true时，first是他与当前节点的距离
        }
        if (l_tag || r_tag) // 包含start节点才能计算他的最长直径
        {
            ans = max(ans, l_len + r_len);
            return {(l_tag ? l_len : r_len) + 1, true}; // 确定感染节点在哪颗子树（不是根节点继续+1）
        }
        return {max(l_len, r_len) + 1, false}; // 子节点的深度
    }

public:
    int amountOfTime(TreeNode *root, int start)
    {
        this->start = start;
        dfs(root);
        return ans;
    }
};